import os
import torch
from pathlib import Path
from PIL import Image

from models.colorization_model import ColorizationGenerator
from models.image_classifier import ImageClassifier

class ModelManager:
    """模型管理器，负责加载和管理不同的着色模型"""
    
    def __init__(self):
        """初始化模型管理器"""
        # 设置设备
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # 设置模型路径
        self.model_paths = {
            'general': os.path.join('models', 'weights', 'flower.pth'),
            'face': os.path.join('models', 'weights', 'face.pth'),
            'flower': os.path.join('models', 'weights', 'flower.pth'),
            'building': os.path.join('models', 'weights', 'flower.pth'),  # 暂时用flower模型替代
            'classifier': os.path.join('models', 'weights', 'image_classifier.pth')
        }
        
        # 初始化模型字典
        self.models = {}
        
        # 初始化分类器
        self.classifier = None
    
    def load_classifier(self):
        """加载图像分类器"""
        if self.classifier is None:
            classifier_path = self.model_paths.get('classifier')
            if os.path.exists(classifier_path):
                try:
                    self.classifier = ImageClassifier(model_path=classifier_path, device=self.device)
                    print(f"成功加载分类器模型: {classifier_path}")
                except Exception as e:
                    print(f"加载分类器模型时出错: {e}")
                    print("使用预训练的基础分类器")
                    self.classifier = ImageClassifier(device=self.device)
            else:
                print(f"分类器模型文件不存在: {classifier_path}")
                print("使用预训练的基础分类器")
                self.classifier = ImageClassifier(device=self.device)
        
        return self.classifier
    
    def _fix_state_dict_keys(self, state_dict):
        """
        修复状态字典中的键名前缀问题
        
        参数:
            state_dict (dict): 原始状态字典
            
        返回:
            dict: 修复后的状态字典
        """
        fixed_dict = {}
        # 检查是否所有键都有"model."前缀
        has_model_prefix = all(k.startswith('model.') for k in state_dict.keys() if k not in ['optimizer_state_dict', 'scheduler_state_dict'])
        
        for key, value in state_dict.items():
            if key in ['optimizer_state_dict', 'scheduler_state_dict']:
                continue
                
            if has_model_prefix and key.startswith('model.'):
                # 移除"model."前缀
                new_key = key[6:]  # 'model.' 长度为6
                fixed_dict[new_key] = value
            elif not has_model_prefix and not key.startswith('model.'):
                # 添加"model."前缀
                fixed_dict[key] = value
            else:
                # 保持原样
                fixed_dict[key] = value
                
        return fixed_dict
    
    def load_model(self, model_type='general'):
        """
        加载指定类型的着色模型
        
        参数:
            model_type (str): 模型类型，'general', 'face', 'flower', 或 'building'
        
        返回:
            model: 加载的模型
        """
        # 如果模型已加载，直接返回
        if model_type in self.models:
            return self.models[model_type]
        
        # 获取模型路径
        model_path = self.model_paths.get(model_type, self.model_paths['general'])
        
        # 检查模型文件是否存在
        if not os.path.exists(model_path):
            print(f"模型文件不存在: {model_path}")
            # 如果指定类型的模型不存在，尝试加载通用模型
            if model_type != 'general' and os.path.exists(self.model_paths['general']):
                print(f"加载通用模型: {self.model_paths['general']}")
                model_path = self.model_paths['general']
                model_type = 'general'
            else:
                print("所有模型文件均不存在，使用未训练的模型")
                model = ColorizationGenerator().to(self.device)
                model.eval()
                self.models[model_type] = model
                return model
        
        # 创建模型实例
        model = ColorizationGenerator().to(self.device)
        
        # 加载模型权重
        try:
            # 使用安全模式加载模型
            checkpoint = torch.load(model_path, map_location=self.device, weights_only=True)
            
            # 处理checkpoint可能是字典的情况
            if isinstance(checkpoint, dict):
                if 'model_state_dict' in checkpoint:
                    state_dict = checkpoint['model_state_dict']
                elif 'state_dict' in checkpoint:
                    state_dict = checkpoint['state_dict']
                else:
                    state_dict = checkpoint
            else:
                state_dict = checkpoint
            
            # 修复状态字典键名
            fixed_state_dict = self._fix_state_dict_keys(state_dict)
            
            # 尝试加载修复后的状态字典
            try:
                model.load_state_dict(fixed_state_dict)
                print(f"模型已成功加载: {model_path}")
            except Exception as e:
                print(f"加载修复后的状态字典失败: {e}")
                print("尝试进行严格=False的加载...")
                # 尝试非严格加载
                model.load_state_dict(fixed_state_dict, strict=False)
                print(f"模型已非严格加载: {model_path}")
                
        except Exception as e:
            print(f"模型加载失败: {e}")
            print("使用未训练的模型")
        
        # 设置为评估模式
        model.eval()
        
        # 缓存模型以备后用
        self.models[model_type] = model
        
        return model
    
    def classify_and_get_model(self, image):
        """
        对图像进行分类并返回相应的着色模型
        
        参数:
            image (PIL.Image): 输入图像
        
        返回:
            model: 适合该图像类型的着色模型
            image_class (str): 图像类别
        """
        # 加载分类器
        classifier = self.load_classifier()
        
        # 分类图像
        image_class = classifier.classify(image)
        print(f"图像分类结果: {image_class}")
        
        # 根据分类结果选择模型
        if image_class == 'face':
            print(f"加载人脸着色模型: {self.model_paths['face']}")
            model = self.load_model('face')
        elif image_class == 'flower':
            print(f"加载花卉着色模型: {self.model_paths['flower']}")
            model = self.load_model('flower')
        else:
            # 这里实际上不会执行到，因为我们的分类器只有两个类别
            print(f"未知分类，使用通用模型: {self.model_paths['general']}")
            model = self.load_model('general')
        
        return model, image_class 